package cn.ycc1.my.rag.controller;

import cn.ycc1.my.rag.service.ChatRagService;
import com.alibaba.dashscope.embeddings.TextEmbedding;
import com.alibaba.dashscope.embeddings.TextEmbeddingParam;
import com.alibaba.dashscope.embeddings.TextEmbeddingResult;
import com.alibaba.dashscope.exception.NoApiKeyException;
import dev.langchain4j.model.chat.ChatLanguageModel;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 对话基于RAG
 * @author ycc
 * @date 2025/4/19
 */
@RestController
@RequestMapping("/chat")
public class ChatController {

    @Autowired
    ChatLanguageModel model;

    private static final String CLUSTER_ENDPOINT = "http://"+ System.getenv("DEEPSEEK_MILVUS_HOST") +":19530";
    private static final String TOKEN = "root:Milvus";

    @Autowired
    private ChatRagService chatRagService;

    @GetMapping("/rag")
    public String chatRag(@RequestParam("prompt") String prompt) throws NoApiKeyException {
        System.out.println("prompt: " + prompt);
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(CLUSTER_ENDPOINT) // 1.2 获取的 Milvus 链接端点
                .token(TOKEN)  // 1.2 获取的 Milvus 链接信息
                .build();

        MilvusClientV2 milvusClientV2 = new MilvusClientV2(connectConfig);

        // List<Float> floatList = embeddingClientOptional.get().embed(prompt);
        // 1. 将用户问题向量化
//        EmbeddingResult queryEmbedding = embedding.call(
//                new EmbeddingParam()
//                        .setModel(TextEmbedding.Models.TEXT_EMBEDDING_V2)
//                        .setTexts(List.of(question)));
        TextEmbeddingParam param = TextEmbeddingParam
                .builder()
                .model(TextEmbedding.Models.TEXT_EMBEDDING_V3)
                .texts(Arrays.asList(prompt)).build();
        TextEmbedding textEmbedding = new TextEmbedding();
        TextEmbeddingResult result2 = textEmbedding.call(param);
        List<Float> floatList = result2.getOutput().getEmbeddings().get(0).getEmbedding()
                .stream()
                .map(Double::floatValue)
                .collect(Collectors.toList());

        SearchReq searchReq = SearchReq.builder()
                .collectionName("deepseek4j_test")
                .data(Collections.singletonList(new FloatVec(floatList)))
                .outputFields(Collections.singletonList("my_varchar"))
                .topK(3)
                .build();

        SearchResp searchResp = milvusClientV2.search(searchReq);

        List<String> resultList = new ArrayList<>();
        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        for (List<SearchResp.SearchResult> results : searchResults) {
            System.out.println("TopK results:");
            for (SearchResp.SearchResult result : results) {
                resultList.add(result.getEntity().get("my_varchar").toString());
            }
        }

        // String responseText = model.generate("你是谁？");
        System.out.println("rag 文本 ===================================");
        System.out.println(resultList);
        String responseText = model.generate(String.format("你要根据用户输入的问题：%s \n \n 参考如下内容： %s  \n\n 整理处理最终结果", prompt, resultList));
        System.out.println(responseText);
        return responseText;
    }

    /**
     * 本地embedding版本RAG
     * @param prompt
     * @return
     */
    @GetMapping("/rag2")
    public String chatRag2(@RequestParam("prompt") String prompt) {
        return chatRagService.localEmbedding(prompt);
//        return "hello";
    }
}
